package com.zixi.ai.rag.service.impl;

import com.alibaba.fastjson2.JSONObject;
import com.zixi.ai.ai.service.IAiModelService;
import com.zixi.ai.ai.service.IVectorStoreService;
import com.zixi.ai.common.constant.Constants;
import com.zixi.ai.common.domain.R;
import com.zixi.ai.common.utils.TextSplitterUtils;
import com.zixi.ai.framework.service.BaseServiceImpl;
import com.zixi.ai.llm.entity.AiLlmEntity;
import com.zixi.ai.llm.service.IAiLlmService;
import com.zixi.ai.rag.dto.DocumentDto;
import com.zixi.ai.rag.entity.AiKnowledgeEntity;
import com.zixi.ai.rag.service.IAiKnowledgeService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.model.transformer.KeywordMetadataEnricher;
import org.springframework.ai.model.transformer.SummaryMetadataEnricher;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.*;

/**
 * @author wfg
 */
@Slf4j
@Service
@Transactional(rollbackFor = Exception.class)
public class AiKnowledgeServiceImpl extends BaseServiceImpl<AiKnowledgeEntity, Long> implements IAiKnowledgeService {
    @Autowired
    private Map<String, IVectorStoreService> vectorStoreServices;
    @Autowired
    private Map<String, IAiModelService> aiModelServices;

    @Autowired
    @Qualifier("qwenChatModel")
    private ChatModel chatModel;
    @Autowired
    private IAiLlmService aiLlmService;
    @Override
    public R<List<Document>> search(Long id, String keyword) {
        VectorStore vs = getVectorStoreById(id);
        List<Document> rs = null;
        if (vs != null) {
            rs = vs.similaritySearch(keyword);
        }
        return R.ok(rs);
    }
    @Override
    public VectorStore getVectorStoreById(Long id) {
        Optional<AiKnowledgeEntity> aks = this.findById(id);
        if (aks.isPresent()) {
            EmbeddingModel edm = getEmbeddingModel(aks.get().getVectorEmbedLlmId());
            if (edm != null) {
                JSONObject storeOptions = aks.get().getOptions();
                if (storeOptions == null) {
                    storeOptions = new JSONObject();
                }
                storeOptions.put("vectorStoreCollection", aks.get().getVectorStoreCollection());
                return vectorStoreServices.get(aks.get().getVectorStoreType() + Constants.VECTOR_STORE_SERVICE)
                        .buildVectorStore(storeOptions, edm, null);
            }
        }
        return null;
    }
    private EmbeddingModel getEmbeddingModel(Long llmId) {
        Optional<AiLlmEntity> llm = aiLlmService.findById(llmId);
        if (llm.isPresent()) {
            JSONObject options = llm.get().getOptions();
            options.put("baseUrl", llm.get().getLlmEndpoint());
            options.put("model", llm.get().getLlmModel());
            return aiModelServices.get(llm.get().getBrand() + Constants.AI_MODEL).buildEmbeddingModel(options);
        }
        return null;
    }
    @Override
    public Boolean saveText(Long knowledgeId, List<DocumentDto> docs) {

        List<Document> documents = new ArrayList<>();
        List<Document> finalDocuments = documents;
        docs.forEach(doc -> {
            if (doc.getMetadata() == null) {
                doc.setMetadata(new HashMap<>());
            }
            if (doc.getMedia() != null) {
                finalDocuments.add(new Document(doc.getMedia(), doc.getMetadata()));
            } else {
                finalDocuments.add(new Document(doc.getText(), doc.getMetadata()));
            }
        });
        log.debug("fileId 写入元数据...");
        documents.forEach(doc -> doc.getMetadata().put("knowledgeId", knowledgeId));

        log.debug("正在提取关键词...");
        KeywordMetadataEnricher keywordEnricher = new KeywordMetadataEnricher(chatModel, 50);
        keywordEnricher.apply(documents);
//
        log.debug("正在生成摘要...");
        SummaryMetadataEnricher summaryEnricher = new SummaryMetadataEnricher(chatModel, List.of(SummaryMetadataEnricher.SummaryType.CURRENT));
        summaryEnricher.apply(documents);

        log.debug("正在拆分文档...");
        documents = splitDocuments(documents);

        log.debug("正在保存文档...");
        VectorStore vs = getVectorStoreById(knowledgeId);
        if (vs != null) {
            vs.add(documents);
        }
        return vs != null;
    }

    /**
     * 分割文档
     */
    public List<Document> splitDocuments(List<Document> documents) {
        List<Document> result = new ArrayList<>();
        for (Document doc : documents) {
            List<String> chunks = TextSplitterUtils.splitText(doc.getText());
            for (String chunk : chunks) {
                Document newDoc = new Document("1", chunk, doc.getMetadata());
                newDoc.getMetadata().putAll(doc.getMetadata());
                result.add(newDoc);
            }
        }
        return result;
    }
}
